from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, IntegerType, StringType

if __name__ == '__main__':
    # 构建SparkSession对象
    spark = SparkSession.builder. \
        appName("local[*]"). \
        config("spark.sql.shuffle.partitions", "4"). \
        getOrCreate()
    # appName 设置程序名称
    # config 设置常用属性。可以通过此来设置配置
    # 最后通过getOrCreate 创建 SparkSession对象

    # 从SparkSession中获取SparkContext
    sc = spark.sparkContext

    """
    需求
    1.查询用户平均分
    2.查询电影平均分
    3.查询大于平均分的电影数量
    4.查询高分电影（>3）打分次数最多的用户，并求此人的打分平均分
    5.查询每个用户的平均打分，最低分，最高分
    6.查询被评分超过100次的电影，的平均分 排名Top10
    """

    # 1.获取数据集
    schema = StructType() \
        .add("user_id", StringType(), nullable=True) \
        .add("movie_id", IntegerType(), nullable=True) \
        .add("rank", IntegerType(), nullable=True) \
        .add("ts", StringType(), nullable=True)

    df = spark.read.format("csv") \
        .option("sep", "\t") \
        .option("header", "false") \
        .option("encoding", "utf-8") \
        .schema(schema) \
        .load("hdfs://node1:8020/input/u.data")

    # 2. 注册为临时表,方便后续操作
    df.createTempView("movie")

    # df.printSchema()
    # df.show()

    # TODO 1:用户平均分
    df.groupBy("user_id") \
        .avg("rank") \
        .withColumnRenamed("avg(rank)", "avg_rank") \
        .withColumn("avg_rank", F.round("avg_rank", 2)) \
        .orderBy("avg_rank", ascending=False) \
        .show()

    # TODO 2:电影平均分
    spark.sql("""
        select movie_id,round(avg(rank),2) as avg_rank  from movie group by movie_id order by avg_rank desc
    """).show()

    # TODO 3:查询大于平均分的电影数量
    # first():获取到Row对象，在通过Row对象获取到某一列的具体值
    avg_rank = df.select(F.avg(df['rank'])).first()['avg(rank)']
    cnt = df.where(df['rank'] > avg_rank).count()
    print(f"大于平均分的电影数量{cnt}")

    # TODO 4:查询高分电影（>3）打分次数最多的用户，并求此人的打分平均分
    # 先找出这个人
    user_id = df.where('rank > 3') \
        .groupBy('user_id') \
        .count() \
        .withColumnRenamed("count", "cnt") \
        .orderBy("cnt", ascending=False) \
        .limit(1) \
        .first()['user_id']

    # 在获取这个人的平均分
    df.filter(df.user_id == user_id).select(F.round(F.avg('rank'), 2)).show()

    # TODO 5:查询每个用户的平均打分，最低分，最高分
    df.groupBy('user_id') \
        .agg(
            F.round(F.avg('rank'), 2).alias('avg_rank'),
            F.round(F.min('rank'), 2).alias('min_rank'),
            F.round(F.max('rank'), 2).alias('max_rank')
        ).show()
    # agg():可以在内写多个聚合函数，且对象为Column

    # TODO 6:查询被评分超过100次的电影，的平均分 排名Top10
    df.groupBy("movie_id") \
        .agg(
            F.round(F.avg('rank'), 2).alias('avg_rank'),
            F.count('movie_id').alias('cnt')
        ).where('cnt > 100') \
        .orderBy('avg_rank', ascending=False) \
        .limit(10).show()
